from pymodelica import compile_fmu
from pyfmi import load_fmu
import matplotlib.pyplot as plt
#from pyjmi.common.plotting import plot_gui

import os as O
import numpy as N
import pylab as p

import sys
import os.path
from time import gmtime, strftime
import argparse

lp = "[TEST SCRIPT] " # Log prefix

def check_env_vars():
    #vars_to_check = ['MODELS', 'MOHOME', 'MODELS', 'MOSCR', 'MOPYSCR', 'sdv', 'vds']
    vars_to_check = ['MODELS', 'MOHOME', 'MODELS', 'MOSCR', 'MOPYSCR']
    vars_not_found = []

    for v in vars_to_check:
        try:
            O.environ[v]
            #print lp + v + " = " + O.environ[v]
        except KeyError:
            #print lp + v + " undefined"
            vars_not_found.append(v)

    return vars_not_found

def get_args():
    parser = argparse.ArgumentParser(description='Simulate a Modelica model.')

    parser.add_argument('-m','--model', 
                        dest='model', 
                        required=True,
                        type=str, 
                        help='Fully qualified model name')

    parser.add_argument('-t', '--time',
                        dest='time',
                        nargs=2, 
                        metavar=('START', 'STOP'),
                        type=float,
                        default=(0.0, 10.0),
                        help='Integration interval')

    parser.add_argument('-s', '--steps',
                        dest='steps',
                        nargs=1, 
                        metavar=('STEPS'),
                        type=int,
                        default=(1000),
                        help='Number of integration steps')

    parser.add_argument('-l', '--sim_log_level',
                        dest='sim_log_level',
                        nargs=1, 
                        metavar=('LEVEL'),
                        type=int,
                        default=(3),
                        help='Simulation log verbosity - 0 (none) to 6 (max).')

    parser.add_argument('-p', '--plot',
                        dest='plot', 
                        required=False,
                        nargs='*',
                        metavar=['<var1>,..,<varN>'],
                        help='Sets of variables to be plotted in separate windows')

    parser.add_argument('-g', '--plot_gui',
                        dest='plot_gui', 
                        required=False,
                        action='store_true',
                        help='Display plot GUI after simulation (does not open the result file)')
 
    parser.add_argument('-n', '--no_comp',
                        dest='no_comp', 
                        required=False,
                        action='store_true',
                        help='Do not recompile model before simulation')
 
    args = parser.parse_args()


    args_dict = {}
    args_dict['model'] = args.model
    args_dict['start'] = args.time[0]
    args_dict['stop']  = args.time[1]
    args_dict['steps'] = args.steps
    args_dict['sim_log_level'] = args.sim_log_level
    if args.plot == None:
        args_dict['plot'] = []
    else:
        args_dict['plot']  = [p.split(',') for p in args.plot]
    args_dict['plot_gui'] = args.plot_gui
    args_dict['no_comp'] = args.no_comp

    print lp + "Arguments: "
    print args_dict
    return args_dict 

def run_test(args):
    # Get Ready ---------------------------------------------------------
    # Check environment variables
    vars_not_found = check_env_vars()
    if not vars_not_found == []:
        raise Exception("Env vars not set: " + ", ".join(vars_not_found))

    # Define paths
    modelica_home_dir = O.environ['MOHOME']
    models_dir = O.environ['MODELS']
    model_file = args['model'].replace('.', '/') + '.mo'
    model_full_name = args['model']
    model_short_name = args['model'][args['model'].rfind('.') + 1:]
    model_name_date = model_short_name + '_' + strftime("%y-%m-%d_%H:%M:%S", gmtime())

    mo_file_path = models_dir + '/' + model_file
    if not os.path.exists(mo_file_path):
        raise Exception("Model file not found: " + mo_file_path)

    print
    print lp + "Model file path: " + mo_file_path
    print lp + "Model short name: " + model_short_name
 
    jmo_ide_dir = modelica_home_dir + "/ide/jmodelica" 
    if not os.path.exists(jmo_ide_dir):
        raise Exception("Compiler output directory not found: " + jmo_ide_dir)
    compiler_log_file = jmo_ide_dir + '/' + model_name_date + '_comp.log'
    simulation_log_file = jmo_ide_dir + '/' + model_name_date + '_simu.log'

    result_dir = modelica_home_dir + '/src/results/' + args['model'] 
    if not os.path.exists(result_dir):
        os.mkdir(result_dir)

    # Generate input
    #t = N.linspace(0., final_time, final_time * 100) 
    #F = 10 * N.cos(t)
    #F_traj = N.transpose(N.vstack((t,F)))

    #input_object = ('F', F_traj)
    
    # Do Work ---------------------------------------------------------
    print lp + '---------------------------------'
 
    # Compile
    if args['no_comp']:
        print lp + '\n\n\nNOT RECOMPILING\n\n\n'
        my_fmu = jmo_ide_dir + '/' + model_full_name.replace('.', '_') + '.fmu' 
        if not os.path.exists(my_fmu):
            raise Exception('FMU not found: ' + my_fmu)
    else:
        print lp + 'COMPILING ' + model_short_name
        #my_fmu = compile_fmu(model_short_name, mo_file_path, 
        my_fmu = compile_fmu(model_full_name, mo_file_path, 
                compile_to=jmo_ide_dir, 
                compiler_log_level='w:' + compiler_log_file)
        print lp + 'Done.'

    print lp + 'Loading FMU ' + my_fmu
    myModel = load_fmu(my_fmu)
    #myModel = load_fmu(my_fmu, log_file_name = simulation_log_file,
                       #log_level = args['sim_log_level'])

    # Set the first input value to the model
    #myModel.set('F', F[0])

    print lp + "Simulating..."
    # Simulate
    res = myModel.simulate(start_time=args['start'], final_time=args['stop'], 
            #input=input_object, \
            options={ 
                'ncp':args['steps'],  
                'result_file_name': result_dir + '/' + model_name_date + '.txt'
            })
    if args['no_comp']:
       print lp + '\n\n\n=================\nDID NOT RECOMPILE\n=================\n\n\n'
 
    return res 


def main():
 
    args = get_args()
    
    try:
        res = run_test(args)
    except Exception as e:
        print lp + "caught exception:"
        print e
        return

    # Plot
    num_win = 1
    for v_win in args['plot']:
        plt.figure(num_win)
        for v in v_win:
            print lp + 'Plotting ' + v
            plt.plot(res['time'], res[v])
        plt.title('Model A.')
        plt.ylabel(', '.join(v_win))
        plt.xlabel('Time (s)')
        plt.legend(v_win)
        num_win = num_win + 1
    plt.show()

    if args['plot_gui']:
        from pyjmi.common.plotting import plot_gui
        plot_gui.startGUI()


main()

